"""
Alibaba Cloud instance management for Dumbo benchmark
"""

import json
import time
from typing import List, Dict, Any
from alibabacloud_ecs20140526.client import Client as Ecs20140526Client
from alibabacloud_ecs20140526 import models as ecs_20140526_models
from alibabacloud_tea_openapi import models as open_api_models
from benchmark.utils import BenchError, Print


class AlibabaInstanceManager:
    """Alibaba Cloud instance manager for Dumbo benchmark"""
    
    def __init__(self):
        self.settings = self._load_settings()
        self.client = self._create_client()
        self.instances = []
    
    def _load_settings(self) -> Dict[str, Any]:
        """Load settings from JSON file"""
        import os
        settings_path = os.path.join(os.path.dirname(__file__), '..', 'settings.json')
        try:
            with open(settings_path, 'r') as f:
                return json.load(f)
        except FileNotFoundError:
            Print.warning(f"Settings file not found: {settings_path}")
            return {}
    
    def _create_client(self):
        """Create Alibaba Cloud ECS client"""
        try:
            # Load credentials
            cred_path = self.settings.get('key', {}).get('accesskey', '/root/.aliyun/access.json')
            with open(cred_path, 'r') as f:
                creds = json.load(f)
            
            config = open_api_models.Config(
                access_key_id=creds.get('AccessKey ID', ''),
                access_key_secret=creds.get('AccessKey Secret', ''),
                endpoint='ecs-cn-hangzhou.aliyuncs.com'
            )
            
            return Ecs20140526Client(config)
        except Exception as e:
            Print.warning(f"Failed to create ECS client: {e}")
            return None
    
    def create_instances(self, nodes: int, instance_type: str = 'ecs.g6e.xlarge', region: str = 'cn-hangzhou') -> List[Dict[str, Any]]:
        """Create Alibaba Cloud instances for Dumbo benchmark"""
        Print.info(f"Creating {nodes} Alibaba Cloud instances of type {instance_type} in {region}")
        
        if not self.client:
            Print.warning("ECS client not available, using mock mode")
            self._mock_create_instances(nodes)
            return self.instances
        
        try:
            # Create instances
            request = ecs_20140526_models.RunInstancesRequest(
                region_id=region,
                instance_type=instance_type,
                image_id='ubuntu_20_04_x64_20G_alibase_20210318.vhd',
                key_pair_name=self.settings.get('key', {}).get('name', 'dumbo-alibaba-key'),
                min_amount=nodes,
                max_amount=nodes,
                security_group_ids=[self._get_security_group_id(region)],
                instance_name=f'dumbo-benchmark-{int(time.time())}'
            )
            
            response = self.client.run_instances(request)
            
            if response.body.instance_id_sets:
                for instance_id in response.body.instance_id_sets.instance_id_set:
                    self.instances.append({
                        'id': instance_id,
                        'type': instance_type,
                        'region': region,
                        'state': 'pending'
                    })
            
            Print.info(f"Created {len(self.instances)} instances")
            
            # Wait for instances to be running
            self._wait_for_instances()
            
            return self.instances
            
        except Exception as e:
            raise BenchError(f"Failed to create Alibaba Cloud instances: {e}")
    
    def _mock_create_instances(self, nodes: int):
        """Create mock instances for testing"""
        for i in range(nodes):
            self.instances.append({
                'id': f'mock-instance-{i}',
                'type': 'ecs.g6e.xlarge',
                'region': 'cn-hangzhou',
                'state': 'running',
                'public_ip': f'127.0.0.{i+1}',
                'private_ip': f'192.168.1.{i+1}'
            })
        Print.info(f"Created {nodes} mock instances")
    
    def _get_security_group_id(self, region: str) -> str:
        """Get security group ID for the region"""
        # This should be configured in settings.json
        return self.settings.get('security_groups', {}).get(region, 'sg-12345678')
    
    def _wait_for_instances(self):
        """Wait for all instances to be running"""
        Print.info("Waiting for instances to be running...")
        
        if not self.client:
            Print.info("Mock mode: instances are already running")
            return
        
        instance_ids = [inst['id'] for inst in self.instances]
        
        # Wait for instances to be running
        max_wait = 300  # 5 minutes
        wait_time = 0
        
        while wait_time < max_wait:
            try:
                request = ecs_20140526_models.DescribeInstancesRequest(
                    instance_ids=','.join(instance_ids)
                )
                response = self.client.describe_instances(request)
                
                all_running = True
                for instance in response.body.instances.instance:
                    state = instance.status
                    if state != 'Running':
                        all_running = False
                        break
                    
                    # Update instance info
                    for i, inst in enumerate(self.instances):
                        if inst['id'] == instance.instance_id:
                            self.instances[i]['public_ip'] = instance.public_ip_address
                            self.instances[i]['private_ip'] = instance.vpc_attributes.private_ip_address.ip_address[0]
                            self.instances[i]['state'] = state
                            break
                
                if all_running:
                    Print.success("All instances are running")
                    return
                
                time.sleep(10)
                wait_time += 10
                
            except Exception as e:
                Print.warning(f"Error checking instance status: {e}")
                time.sleep(10)
                wait_time += 10
        
        raise BenchError("Timeout waiting for instances to be running")
    
    def get_instance_info(self) -> List[Dict[str, Any]]:
        """Get information about all instances"""
        return self.instances
    
    def terminate_instances(self):
        """Terminate all instances"""
        if not self.instances:
            Print.info("No instances to terminate")
            return
        
        Print.info("Terminating Alibaba Cloud instances...")
        
        if not self.client:
            Print.info("Mock mode: no actual termination needed")
            self.instances = []
            return
        
        instance_ids = [inst['id'] for inst in self.instances]
        
        try:
            request = ecs_20140526_models.DeleteInstancesRequest(
                instance_id=instance_ids,
                force=True
            )
            
            self.client.delete_instances(request)
            Print.success("All instances terminated")
            
            # Clear instances list
            self.instances = []
            
        except Exception as e:
            raise BenchError(f"Failed to terminate instances: {e}")
    
    def create_security_group(self, group_name: str = 'dumbo-benchmark-sg') -> str:
        """Create security group for Dumbo benchmark"""
        Print.info(f"Creating security group: {group_name}")
        
        if not self.client:
            Print.info("Mock mode: returning mock security group ID")
            return 'sg-mock-12345678'
        
        try:
            # Create security group
            request = ecs_20140526_models.CreateSecurityGroupRequest(
                region_id='cn-hangzhou',
                security_group_name=group_name,
                description='Security group for Dumbo benchmark'
            )
            
            response = self.client.create_security_group(request)
            group_id = response.body.security_group_id
            
            # Add inbound rules
            authorize_request = ecs_20140526_models.AuthorizeSecurityGroupRequest(
                region_id='cn-hangzhou',
                security_group_id=group_id,
                ip_protocol='tcp',
                port_range='22/22',
                source_cidr_ip='0.0.0.0/0'
            )
            self.client.authorize_security_group(authorize_request)
            
            authorize_request2 = ecs_20140526_models.AuthorizeSecurityGroupRequest(
                region_id='cn-hangzhou',
                security_group_id=group_id,
                ip_protocol='tcp',
                port_range='9000/9100',
                source_cidr_ip='0.0.0.0/0'
            )
            self.client.authorize_security_group(authorize_request2)
            
            Print.success(f"Security group created: {group_id}")
            return group_id
            
        except Exception as e:
            raise BenchError(f"Failed to create security group: {e}")
    
    def create_key_pair(self, key_name: str = 'dumbo-alibaba-key') -> str:
        """Create key pair for Alibaba Cloud instances"""
        Print.info(f"Creating key pair: {key_name}")
        
        if not self.client:
            Print.info("Mock mode: creating mock key file")
            key_file = f"{key_name}.pem"
            with open(key_file, 'w') as f:
                f.write("-----BEGIN RSA PRIVATE KEY-----\nMOCK_KEY\n-----END RSA PRIVATE KEY-----")
            return key_file
        
        try:
            request = ecs_20140526_models.CreateKeyPairRequest(
                region_id='cn-hangzhou',
                key_pair_name=key_name
            )
            
            response = self.client.create_key_pair(request)
            
            # Save private key to file
            key_file = f"{key_name}.pem"
            with open(key_file, 'w') as f:
                f.write(response.body.private_key_body)
            
            # Set proper permissions
            import os
            os.chmod(key_file, 0o400)
            
            Print.success(f"Key pair created: {key_file}")
            return key_file
            
        except Exception as e:
            raise BenchError(f"Failed to create key pair: {e}")

